# -*- coding:utf-8 -*-

# @Time    : 2023/5/16 02:21
# @Author  : zengwenjia
# @Email   : zengwenjia@lingxi.ai
# @File    : generate_bot_dialogue.py
# @Software: LLM_internal

# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
import pandas as pd
import asyncio
import uuid
from data_generate import utils
import traceback
from agent.llm_agent import LLMAgent
from bot.financial_sales.sales import Sales
import sys
import random
import os
import tqdm
import jsonlines

model_name = "Colossal-7B-financial-1019"

def read_prompt_data(path):
    print(path)
    new_data = []
    data = utils.jload(path)
    for instruct_dict in tqdm.tqdm(data):
        conversations = instruct_dict.get("conversations", "")
        for k in conversations:
            if k['from'] == 'human':
                prompt = k['value']
            if k['from'] == 'gpt':
                gpt_response = k['value']
            chat = LLMAgent(prompt)
            self_model_response = asyncio.run(chat.achat_auto_llm(type=model_name))
        if gpt_response != self_model_response:
            chosen_data = "\n\nHuman: " + prompt.strip() + "\n\nAssistant: " + gpt_response.strip()
            reject_data = "\n\nHuman: " + prompt.strip() + "\n\nAssistant: " + self_model_response.strip()
            new_data.append({
                "chosen:": chosen_data,
                "reject:": reject_data,
            })

    return new_data


async def read_dialogue_data(path):
    print(path)
    new_data = []
    data = utils.jload(path)
    bot = Sales()
    for instruct_dict in tqdm.tqdm(data):
        session_id = instruct_dict.get("id", "")
        messages = instruct_dict.get("messages", "")
        context = []
        base_info = []
        for i in range(len(messages)):
            context.append(messages[i])
            if i < 2:
                if i == 1:
                    if "女士" in messages[i]['content']:
                        index = messages[i]['content'].find("女士")
                    if "先生" in messages[i]['content']:
                        index = messages[i]['content'].find("先生")
                    name = messages[i]['content'][index - 1:index + 2]
                continue
            if messages[i]['role'] == 'user':
                res = await bot.async_reply(name, context, session_id, base_info)
            if messages[i]['role'] == 'assistant':
                content = messages[i]['content']
                if content != res:
                    j = 0
                    chosen_data = ""
                    reject_data = ""
                    while j < len(context) - 1:
                        if context[j]['role'] == 'user':
                            chosen_data += "\n\nHuman: {}".format(context[j]['content'].strip())
                            reject_data += "\n\nHuman: {}".format(context[j]['content'].strip())
                        if context[j]['role'] == 'assistant':
                            chosen_data += "\n\nAssistant: {}".format(context[j]['content'].strip())
                            reject_data += "\n\nAssistant: {}".format(context[j]['content'].strip())
                        j += 1
                    if j == len(context) - 1:
                        chosen_data += "\n\nAssistant: {}".format(content)
                        reject_data += "\n\nAssistant: {}".format(res)
                        new_data.append({
                            "chosen:": chosen_data,
                            "reject:": reject_data,
                        })

    return new_data


def generate_rm_data():
    # prompt_data_1 = read_prompt_data("../internal_server/train_data/llm/2023-10-16.json")
    # prompt_data_2 = read_prompt_data("../internal_server/train_data/llm/2023-10-17.json")
    prompt_data_3 = read_prompt_data("../internal_server/train_data/llm/2023-10-18.json")
    prompt_data_4 = read_prompt_data("../internal_server/train_data/llm/2023-10-19.json")
    # dialogue_data_1 = asyncio.run(read_dialogue_data("../data_set/bot_dialogue/bot_dialogue_2023-10-16.json"))
    # dialogue_data_2 = asyncio.run(read_dialogue_data("../data_set/bot_dialogue/bot_dialogue_2023-10-17.json"))
    dialogue_data_3 = asyncio.run(read_dialogue_data("../data_set/bot_dialogue/bot_dialogue_2023-10-18.json"))
    dialogue_data_4 = asyncio.run(read_dialogue_data("../data_set/bot_dialogue/bot_dialogue_2023-10-19.json"))

    d = prompt_data_3 + prompt_data_4 + dialogue_data_3 + dialogue_data_4

    with jsonlines.open("rm_data_1023.jsonl", "w") as writer:
        writer.write_all(d)


if __name__ == '__main__':
    generate_rm_data()
